package adkr

import (
	"ebbflow/pkg/crypto"
	"encoding/json"
	"fmt"
	"log"
	"sync"
	"time"
)

// ADKRHighConfig represents the configuration for ADKR high-level protocol
type ADKRHighConfig struct {
	SID        string        `json:"sid"`
	PID        int           `json:"pid"`
	N          int           `json:"n"`
	F          int           `json:"f"`
	K          int           `json:"k"` // Number of rounds
	G1         *crypto.G1    `json:"g1"`
	G2         *crypto.G2    `json:"g2"`
	EPKs       []crypto.G1   `json:"epks"`
	ESK        crypto.ZR     `json:"esk"`
	THPK       crypto.G1     `json:"thpk"`
	THPKs      []crypto.G1   `json:"thpks"`
	THSK       crypto.ZR     `json:"thsk"`
	Mute       bool          `json:"mute"`
	Debug      bool          `json:"debug"`
}

// ADKRHighOutput represents the output of ADKR high-level protocol
type ADKRHighOutput struct {
	NewTHPK     crypto.G1     `json:"new_thpk"`
	NewTHPKs    []crypto.G1   `json:"new_thpks"`
	NewTHSK     crypto.ZR     `json:"new_thsk"`
	ConfigChain [][]int       `json:"config_chain"`
	Statistics  ProtocolStats `json:"statistics"`
}

// ProtocolStats represents protocol execution statistics
type ProtocolStats struct {
	StartTime    time.Time `json:"start_time"`
	EndTime      time.Time `json:"end_time"`
	Duration     float64   `json:"duration_seconds"`
	RoundsRun    int       `json:"rounds_run"`
	MessagesSent int       `json:"messages_sent"`
	MessagesRecv int       `json:"messages_received"`
	NetworkIO    NetworkIO `json:"network_io"`
}

// NetworkIO represents network I/O statistics
type NetworkIO struct {
	BytesSent    int64 `json:"bytes_sent"`
	BytesRecv    int64 `json:"bytes_received"`
	PacketsSent  int64 `json:"packets_sent"`
	PacketsRecv  int64 `json:"packets_received"`
}

// ADKRHighMessage represents messages in ADKR high-level protocol
type ADKRHighMessage struct {
	Tag     string `json:"tag"`
	Round   int    `json:"round"`
	Payload []byte `json:"payload"`
}

// ADKRHighTag represents different message types in ADKR high-level protocol
type ADKRHighTag string

const (
	ADKR_NEW   ADKRHighTag = "ADKR_NEW"
	ADKR_OLD   ADKRHighTag = "ADKR_OLD"
	ADKR_COUNT ADKRHighTag = "ADKR_COUNT"
)

// ADKRHighProtocol implements the high-level ADKR protocol
type ADKRHighProtocol struct {
	config           ADKRHighConfig
	msgInCh          <-chan Message
	msgOutCh         chan<- Message
	outputCh         chan<- ADKRHighOutput
	logger           *log.Logger
	mu               sync.RWMutex
	stats            ProtocolStats
	roundResults     []KeyRefreshOutput
	configChain      [][]int
	currentConfig    []int
	nextConfig       []int
	completedRounds  int
	messageCount     int
	networkStats     NetworkIO
}

// NewADKRHighProtocol creates a new ADKR high-level protocol instance
func NewADKRHighProtocol(
	config ADKRHighConfig,
	msgInCh <-chan Message,
	msgOutCh chan<- Message,
	outputCh chan<- ADKRHighOutput,
	logger *log.Logger,
) *ADKRHighProtocol {
	// Initialize configuration chain
	configChain := make([][]int, config.K+1)
	for i := 0; i <= config.K; i++ {
		configChain[i] = make([]int, config.N)
		for j := 0; j < config.N; j++ {
			configChain[i][j] = j
		}
	}
	
	return &ADKRHighProtocol{
		config:        config,
		msgInCh:       msgInCh,
		msgOutCh:      msgOutCh,
		outputCh:      outputCh,
		logger:        logger,
		stats:         ProtocolStats{StartTime: time.Now()},
		roundResults:  make([]KeyRefreshOutput, 0, config.K),
		configChain:   configChain,
		currentConfig: configChain[0],
		nextConfig:    configChain[1],
	}
}

// Run executes the ADKR high-level protocol
func (ahp *ADKRHighProtocol) Run() error {
	ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Starting ADKR high-level protocol", ahp.config.SID, ahp.config.PID)
	
	// Start message processing
	go ahp.processMessages()
	
	// Run K rounds of key refresh
	for round := 0; round < ahp.config.K; round++ {
		ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Starting round %d", ahp.config.SID, ahp.config.PID, round)
		
		roundStart := time.Now()
		
		// Run key refresh for this round
		result, err := ahp.runKeyRefreshRound(round)
		if err != nil {
			return fmt.Errorf("round %d failed: %w", round, err)
		}
		
		// Store round result
		ahp.mu.Lock()
		ahp.roundResults = append(ahp.roundResults, *result)
		ahp.completedRounds++
		ahp.mu.Unlock()
		
		roundDuration := time.Since(roundStart)
		ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Completed round %d in %v", 
			ahp.config.SID, ahp.config.PID, round, roundDuration)
		
		// Update configuration for next round
		ahp.updateConfigForNextRound(round, result)
	}
	
	// Generate final output
	finalOutput := ahp.generateFinalOutput()
	ahp.outputCh <- finalOutput
	
	ahp.logger.Printf("ADKRHigh-SID:%s PID:%d ADKR high-level protocol completed", ahp.config.SID, ahp.config.PID)
	return nil
}

// runKeyRefreshRound runs a single round of key refresh
func (ahp *ADKRHighProtocol) runKeyRefreshRound(round int) (*KeyRefreshOutput, error) {
	ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Running key refresh round %d", ahp.config.SID, ahp.config.PID, round)
	
	// Create key refresh configuration
	keyRefreshConfig := KeyRefreshConfig{
		SID:        ahp.config.SID + fmt.Sprintf("-round-%d", round),
		PID:        ahp.config.PID,
		N:          ahp.config.N,
		F:          ahp.config.F,
		L:          ahp.config.F, // Assuming l = f
		ConfigChain: [][]int{ahp.currentConfig, ahp.nextConfig},
		G:          ahp.config.G1,
		Type:       "b", // BN254 type
		SPK2s:      ahp.config.THPKs,
		SSK2:       ahp.config.THSK,
		EPKs:       ahp.config.EPKs,
		ESK:        ahp.config.ESK,
		THSKOld:    ahp.config.THSK,
		THPKOld:    ahp.config.THPK,
		THPKsOld:   ahp.config.THPKs,
	}
	
	// Create channels for key refresh
	keyRefreshInCh := make(chan Message, 100)
	keyRefreshOutCh := make(chan Message, 100)
	keyRefreshOutputCh := make(chan KeyRefreshOutput, 1)
	
	// Create key refresh protocol
	keyRefreshProtocol := NewKeyRefreshProtocol(
		keyRefreshConfig,
		keyRefreshInCh,
		keyRefreshOutCh,
		keyRefreshOutputCh,
		ahp.logger,
	)
	
	// Start key refresh protocol
	go func() {
		if err := keyRefreshProtocol.Run(); err != nil {
			ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Key refresh round %d failed: %v", 
				ahp.config.SID, ahp.config.PID, round, err)
		}
	}()
	
	// Forward messages between main protocol and key refresh
	go ahp.forwardMessages(keyRefreshInCh, keyRefreshOutCh, round)
	
	// Wait for key refresh to complete
	select {
	case result := <-keyRefreshOutputCh:
		ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Key refresh round %d completed successfully", 
			ahp.config.SID, ahp.config.PID, round)
		return &result, nil
	case <-time.After(func() time.Duration {
		if IsFastBench() { return 45 * time.Second }
		return 90 * time.Second
	}()): // Timeout reduced in FastBench mode
		return nil, fmt.Errorf("key refresh round %d timed out", round)
	}
}

// forwardMessages forwards messages between main protocol and key refresh
func (ahp *ADKRHighProtocol) forwardMessages(keyRefreshInCh chan<- Message, keyRefreshOutCh <-chan Message, round int) {
	// Forward outgoing messages from key refresh to main protocol
	go func() {
		for msg := range keyRefreshOutCh {
			// Add round information and sender to message
			msg.Round = round
			msg.Sender = ahp.config.PID
			ahp.msgOutCh <- msg
			ahp.mu.Lock()
			ahp.messageCount++
			ahp.mu.Unlock()
		}
	}()
	
	// Forward incoming messages from main protocol to key refresh
	go func() {
		for msg := range ahp.msgInCh {
			// Filter messages for this round
			if msg.Round == round {
				select {
				case keyRefreshInCh <- msg:
				case <-time.After(1 * time.Second):
					ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Timeout forwarding message to key refresh", 
						ahp.config.SID, ahp.config.PID)
				}
			}
		}
	}()
}

// processMessages processes incoming messages
func (ahp *ADKRHighProtocol) processMessages() {
	for msg := range ahp.msgInCh {
		ahp.mu.Lock()
		ahp.messageCount++
		ahp.mu.Unlock()
		
		// Process different message types
		switch msg.Tag {
		case BroadcastTag(ADKR_COUNT):
			ahp.handleCountMessage(msg)
		default:
			// Other messages are handled by the key refresh protocol
			ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Received message with tag %s", 
				ahp.config.SID, ahp.config.PID, msg.Tag)
		}
	}
}

// handleCountMessage handles count messages for statistics
func (ahp *ADKRHighProtocol) handleCountMessage(msg Message) {
	var countData struct {
		Time     float64 `json:"time"`
		Network  [2]int64 `json:"network"`
	}
	
	if err := json.Unmarshal(msg.Payload, &countData); err != nil {
		ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Failed to unmarshal count message: %v", 
			ahp.config.SID, ahp.config.PID, err)
		return
	}
	
	ahp.logger.Printf("ADKRHigh-SID:%s PID:%d Node %d: time=%f, network=%v", 
		ahp.config.SID, ahp.config.PID, msg.Sender, countData.Time, countData.Network)
}

// updateConfigForNextRound updates configuration for the next round
func (ahp *ADKRHighProtocol) updateConfigForNextRound(round int, result *KeyRefreshOutput) {
	ahp.mu.Lock()
	defer ahp.mu.Unlock()
	
	// Update threshold keys
	ahp.config.THPK = result.NewTHPK
	ahp.config.THPKs = result.NewTHPKs
	ahp.config.THSK = result.NewTHSK
	
	// Update configuration chain
	if round+1 < len(ahp.configChain) {
		ahp.currentConfig = ahp.configChain[round+1]
		if round+2 < len(ahp.configChain) {
			ahp.nextConfig = ahp.configChain[round+2]
		}
	}
}

// generateFinalOutput generates the final output of the protocol
func (ahp *ADKRHighProtocol) generateFinalOutput() ADKRHighOutput {
	ahp.mu.Lock()
	defer ahp.mu.Unlock()
	
	// Update final statistics
	ahp.stats.EndTime = time.Now()
	ahp.stats.Duration = ahp.stats.EndTime.Sub(ahp.stats.StartTime).Seconds()
	ahp.stats.RoundsRun = ahp.completedRounds
	ahp.stats.MessagesSent = ahp.messageCount
	ahp.stats.MessagesRecv = ahp.messageCount
	
	return ADKRHighOutput{
		NewTHPK:     ahp.config.THPK,
		NewTHPKs:    ahp.config.THPKs,
		NewTHSK:     ahp.config.THSK,
		ConfigChain: ahp.configChain,
		Statistics:  ahp.stats,
	}
}

// GetStatistics returns current protocol statistics
func (ahp *ADKRHighProtocol) GetStatistics() ProtocolStats {
	ahp.mu.RLock()
	defer ahp.mu.RUnlock()
	
	stats := ahp.stats
	stats.EndTime = time.Now()
	stats.Duration = stats.EndTime.Sub(stats.StartTime).Seconds()
	stats.RoundsRun = ahp.completedRounds
	stats.MessagesSent = ahp.messageCount
	stats.MessagesRecv = ahp.messageCount
	
	return stats
}

// ADKRHigh is the main function that orchestrates the high-level ADKR protocol
// This corresponds to the Python class Adkgadphigh
func ADKRHigh(
	sid string, pid, n, f, k int, g1 *crypto.G1, g2 *crypto.G2,
	ePKs []crypto.G1, eSK crypto.ZR, thpk crypto.G1, thpks []crypto.G1, thsk crypto.ZR,
	msgInCh <-chan Message, msgOutCh chan<- Message,
	mute bool, debug bool,
) {
	logger := log.Default()
	logger.Printf("ADKRHigh-SID:%s PID:%d Starting high-level ADKR protocol", sid, pid)
	
    // Paillier keys should be registered via config before running (bench/main registers for demo)

	// Create configuration
	config := ADKRHighConfig{
		SID:   sid,
		PID:   pid,
		N:     n,
		F:     f,
		K:     k,
		G1:    g1,
		G2:    g2,
		EPKs:  ePKs,
		ESK:   eSK,
		THPK:  thpk,
		THPKs: thpks,
		THSK:  thsk,
		Mute:  mute,
		Debug: debug,
	}
	
	// Create output channel
	outputCh := make(chan ADKRHighOutput, 1)
	
	// Create protocol instance
	protocol := NewADKRHighProtocol(
		config,
		msgInCh,
		msgOutCh,
		outputCh,
		logger,
	)
	
	// Run the protocol
	if err := protocol.Run(); err != nil {
		logger.Printf("ADKRHigh-SID:%s PID:%d Protocol failed: %v", sid, pid, err)
		return
	}
	
	// Wait for output
	select {
	case output := <-outputCh:
		logger.Printf("ADKRHigh-SID:%s PID:%d Protocol completed successfully", sid, pid)
		logger.Printf("ADKRHigh-SID:%s PID:%d Final statistics: %+v", sid, pid, output.Statistics)
	case <-time.After(func() time.Duration {
		if IsFastBench() { return 45 * time.Second }
		return 60 * time.Second
	}()):
		logger.Printf("ADKRHigh-SID:%s PID:%d Protocol timed out waiting for output", sid, pid)
	}
}
